# -*- coding: utf-8 -*-

import os
import json
import oyaml as yaml
import numpy as np


class ODCheckBackend(object):
    def __init__(self, yaml_file_path):
        self._pairs = {"pairs": list(), "pairs_index": 0}
        self._yaml_file_path = yaml_file_path

        self.read_yaml_file(yaml_file_path)

    @staticmethod
    def read_text_label_file(label_file_path):
        with open(label_file_path, "rb") as label_file:
            data = np.frombuffer(label_file.read(), dtype=np.int32)
        x1, x2, y1, y2 = data.tolist()
        w = x2 - x1 + 1
        h = y2 - y1 + 1
        return x1, y1, w, h

    @staticmethod
    def read_json_label_file(label_file_path):
        with open(label_file_path, "r") as label_file:
            data = json.load(label_file)
        image_width = data["image"]["information"]["width"]
        image_height = data["image"]["information"]["height"]
        x1, y1, x2, y2 = [float(e) for e in data["annotation"]["persons"][0]["box"]]
        x1 = int(x1 * image_width)
        y1 = int(y1 * image_height)
        x2 = int(x2 * image_width)
        y2 = int(y2 * image_height)
        w = x2 - x1 + 1
        h = y2 - y1 + 1
        return x1, y1, w, h

    def read_label_file(self, image_path):
        label_file_path = image_path.replace(".jpg", ".txt")
        if os.path.exists(label_file_path):
            return self.read_text_label_file(label_file_path)
        label_file_path = image_path.replace(".jpg", ".json")
        if os.path.exists(label_file_path):
            return self.read_json_label_file(label_file_path)
        raise ValueError

    def read_image_and_label_paths(self, folder):
        image_names = [name for name in os.listdir(folder) if ".jpg" in name]
        image_paths = [os.path.join(folder, name) for name in image_names]
        rectangles = [self.read_label_file(image_path) for image_path in image_paths]
        return image_paths, rectangles

    def read_yaml_file(self, yaml_file_path):
        with open(yaml_file_path, "r", encoding="utf-8") as yaml_file:
            data = yaml.load(yaml_file)
        root_folder = data["root"]
        yaml_pairs = data["pairs"]

        # 为了显示的方便，记录下所有图片的完整路径，以及每张图片对应的矩形框位置
        # 同时，保存每个gallery和query当前的序号位置
        for yaml_pair_name, yaml_pair in yaml_pairs.items():
            pair = dict()
            pair["name"] = yaml_pair_name
            # gallery
            gallery_image_folder = os.path.join(root_folder, yaml_pair["gallery"]["image"])
            image_paths, rectangles = self.read_image_and_label_paths(gallery_image_folder)
            pair["gallery"] = {"image": image_paths,
                               "rectangle": rectangles,
                               "time": yaml_pair["gallery"]["time"],
                               "location": yaml_pair["gallery"]["location"],
                               "index": 0}
            # query
            query_image_folder = os.path.join(root_folder, yaml_pair["query"]["image"])
            image_paths, rectangles = self.read_image_and_label_paths(query_image_folder)
            pair["query"] = {"image": image_paths,
                             "rectangle": rectangles,
                             "time": yaml_pair["query"]["time"],
                             "location": yaml_pair["query"]["location"],
                             "index": 0}
            pair["distance"] = yaml_pair["distance"]
            pair["source"] = yaml_pair["source"]
            pair["status"] = yaml_pair["status"]
            self._pairs["pairs"].append(pair)

    def write_checked_result(self):
        with open(self._yaml_file_path, "r", encoding="utf-8") as yaml_file:
            data = yaml.load(yaml_file)
        root_folder = data["root"]
        origin_pairs = data["pairs"]
        wrong_pair_names = list()  # 记录需要删除的pair记录
        for name in origin_pairs:  # 修改每一个pair的status
            for checked_pair in self._pairs["pairs"]:  # self._pairs["pairs"]记录了每个pair当前的状态
                if name == checked_pair["name"]:
                    origin_pairs[name]["status"] = checked_pair["status"]
                    if checked_pair["status"] == "wrong":
                        wrong_pair_names.append(name)
                    break
        for name in wrong_pair_names:  # 去除标记为wrong的记录
            origin_pairs.pop(name)

        with open(self._yaml_file_path, "w", encoding="utf-8") as yaml_file:
            yaml.dump(
                {"root": root_folder, "pairs": origin_pairs},
                yaml_file,
                allow_unicode=True,
                default_flow_style=False)

    @property
    def pairs(self):
        return self._pairs["pairs"]

    @property
    def pairs_index(self):
        return self._pairs["pairs_index"]

    @pairs_index.setter
    def pairs_index(self, index):
        self._pairs["pairs_index"] = index
